Skip to content

Don't allocate when cooperative_wait waits on a worker thread - #39

Merged
maleadt merged 1 commit into
mainfrom
sd/pool-worker-wait-requests
Nov 5, 2026
Merged

maleadt merged 1 commit into
mainfrom
sd/pool-worker-wait-requests

Conversation

@SimonDanisch

@SimonDanisch SimonDanisch commented Nov 3, 2026 •

Copy link
Copy Markdown
Member

Waits that cooperative_wait hands to a worker thread (CUDA.jl, AMDGPU.jl, OpenCL on GPUs) allocate a request with its event, and box the object and the result: 224 to 288 bytes per wait. Whether a wait takes that path depends on how long the operation runs and how loaded the host is, so allocation counts of e.g. a kernel launch followed by synchronize() aren't deterministic. #32 fixed the same problem for waits on completion notifications (PoCL).

Requests are now reused from pools, one per type of wait function, object and inferred result type (promote_op), so the object and the result are stored without boxing. They are kept in 1-element vectors that are emptied on release. That stores immutable objects containing references inline, like CUDA.jl's (stream, context) tuple, which a Union{Nothing,T} field would box. It also keeps idle requests from retaining the driver's objects. Requests of interrupted or cancelled waits are dropped rather than reused, since the worker may still be using them.

The worker now calls wait in the caller's world rather than the latest one, so it sees the same methods as a direct call (and as the in-finalizer path, which already calls wait directly). That keeps the result of the inferred type, and the return type at Union{Nothing,Some{R}}. Supporting a latest-world result of another type would need a Some{Any} in the return type, which boxes every returned result when cooperative_wait isn't inlined. The only behavior change: a method defined after the caller's world, e.g. a redefinition while the wait is in progress, is no longer picked up.

The pools have their own spin lock. Sharing the worker lock cost about 0.6 µs per round trip, since the worker takes that lock just before waking the caller. An interrupt while spinning on the pool lock is deferred, like any other interrupt during the wait.

Per wait on a worker thread (Linux, Julia 1.12):

object, result before after
mutable, nothing 224 B 0 B
mutable, Int32-like 256 B 0 B
(mutable, mutable) tuple, Int32-like 288 B 0 B

With CUDA.jl on an RTX 5080, synchronize() after a long kernel goes from 672 to 400 bytes. The rest is allocated by CUDA.jl itself, mostly when its worker function sets up the context. A worker round trip takes about as long as before (1.0–1.2 µs, against 0.9–1.1). Tested on Julia 1.10–1.14 on Linux, and 1.10 and 1.13 on macOS.

Bumps the version to 3.5.1.

Waits handed to a worker thread (CUDA.jl, AMDGPU.jl, OpenCL on GPUs) allocated a request
with its event, and boxed the object and the result: 224 to 288 bytes per wait. Whether a
wait takes that path depends on the operation's duration and on host load, so allocation
counts weren't deterministic, as with completion notifications before #32.

Requests are now reused, from pools per type of function, object and inferred result, so
that those are stored unboxed. The arguments and the result are kept in vectors that are
emptied on release, which also store immutable objects with references (like CUDA.jl's
`(stream, context)`) inline, and keep idle requests from retaining them. Requests of
interrupted or cancelled waits are not reused.

The worker now calls the function in the caller's world instead of the latest one, as a
direct call would, so that its result always has the inferred type. Otherwise the return
type would need to include another, boxed, `Some`, which makes returning the result
allocate when `cooperative_wait` isn't inlined.

The pools use a spin lock of their own, which defers interrupts like the rest of the wait. A
worker round trip takes about as long as before (1.1-1.3 µs on Linux, Julia 1.12).

Bumps the version to 3.5.1.
@maleadt
maleadt force-pushed the sd/pool-worker-wait-requests branch from f8a5424 to 0de7b0d Compare November 5, 2026 08:48
@maleadt maleadt changed the title Improve allocations and error handling for wait request Don't allocate when cooperative_wait waits on a worker thread Nov 5, 2026
@maleadt
maleadt merged commit 14f1a12 into main Nov 5, 2026
8 of 10 checks passed
@maleadt
maleadt deleted the sd/pool-worker-wait-requests branch November 5, 2026 09:44
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants